#Script to create a velocity map (moment 1 map) from a known "good" line, in
#this case one of the stacked CH3OH transitions, then use that velocity map to
#shift-and-stack all of the spectra in the cube.
#The main functionality is spectral-cube's stacking function:
#https://github.com/radio-astro-tools/spectral-cube/blob/master/spectral_cube/analysis_utilities.py#L136
#Adapted from script from Ginsburg et al. (2018)

from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
import numpy as np
import os
import spectral_cube.analysis_utilities
from spectral_cube import SpectralCube
from astropy import units as u
from astropy.io import fits
import pylab as pl
import regions
import reproject
import pickle
from  radio_beam import *

spectral_dict={}

spectral_dict['H218O_203_template']={'filename': 'H218O-i38.3-203GHz-edited-header.fits'}

radecstring='05h38m18.100454s -07d02m25.99340s'
coord=SkyCoord([radecstring],frame='icrs',unit=(u.hourangle,u.deg))

# step 1: create a velocity map

vmap_name = 'disk_velocity_map_h218o_template.fits'
if not os.path.exists(vmap_name):
    cube = SpectralCube.read('H218O-i38.3-203GHz-edited-header.fits')
    cube2 = cube.with_spectral_unit(u.km / u.s, velocity_convention='radio',rest_value=203.40752 * u.GHz) 
    m1 = cube2.moment1()
    m0 = cube2.moment0()
    mask = m0.value > 0.00001

    vmap = m1
    vmap[~mask] = np.nan

    hdu = vmap.hdu
    #hdu.data = vmap_
    hdu.writeto(vmap_name, overwrite=True)
else:
    hdu = fits.open(vmap_name)[0]
vmap = spectral_cube.lower_dimensional_structures.Projection.from_hdu(hdu)


# step 2: stack

for key in spectral_dict.keys():
   filename=spectral_dict[key]['filename']

   image=fits.open(filename)
   header=image[0].header
   radius=0.4
   radius_px=0.4/(header['CDELT2']*3600.0)
   #npix_beam=3.14159*header['BMAJ']*header['BMIN']/(4.0*np.log(2.0))/header['CDELT2']**2
   #npix=radius_px**2*3.14
   #nbeam_per_spectrum=3.14159*radius_px**2/npix_beam


   # load the cube
   fullcube = (SpectralCube.read(filename, use_dask=False))



   # convert the cube to velocity units with an arbitrary reference point
   # (this step assumes the cube is in frequency or wavelength; if the
   # cube is not, it should be skipped)
   fullcube = fullcube.with_spectral_unit(u.km/u.s,
                                          velocity_convention='radio',
                                          rest_value=203.40752 * u.GHz)
   fullcube.allow_huge_operations=True
   # mask out super bright SiO masers; they break the FFT shifting tool
   # (this step can be skipped if there's nothing anomalously bright
   # in your spectrum)
   #fullcube = fullcube.mask_out_bad_beams(0.1)
   #cb = fullcube.beams.common_beam()
   #fullcube.beam=Beam(0.4*u.arcsec)
   #fullcube = fullcube.to(u.Jy/u.pix)*npix
   fullcube.beam=Beam(0.1*u.arcsec)

   # reproject the velocity map into the cube's coordinate system
   vmap_proj,_ = reproject.reproject_interp(vmap.hdu,
                                            fullcube.wcs.celestial,
                                            shape_out=fullcube.shape[1:])
   vmap_proj = u.Quantity(vmap_proj, u.km/u.s)

   # perform the stacking!
   stack = spectral_cube.analysis_utilities.stack_spectra(fullcube, vmap_proj,
                                                          v0=0.0*u.km/u.s)
   fstack = stack.with_spectral_unit(u.GHz)

   fstack.write('stacked_'+filename,
                overwrite=True)
   stack.write('stacked_cube_'+filename,
                overwrite=True)

   pl.clf()
   fstack.quicklook(filename='stacked_'+filename.replace('.fits','.png'))
   spectral_dict[key]['spectrum']=fstack.value
   spectral_dict[key]['spectrum_Jy_beam']=fstack.value
   spectral_dict[key]['freq_axis']=fstack.spectral_axis.value





pickle.dump(spectral_dict, open("spectral_dict_lime_stacked.pickle", "wb"))

spectral_dict=pickle.load(open("spectral_dict_lime_stacked.pickle", "rb"))
vlsr=0.0
for key in list(spectral_dict.keys()):
   spectral_dict[key]['freq_axis']=spectral_dict[key]['freq_axis']+vlsr/3.0e5*np.median(spectral_dict[key]['freq_axis'])
   outfile=open(key+'_stacked_text_spectra.txt','w')
   for i in range(len(spectral_dict[key]['freq_axis'])):
      line='{} {} \n'.format(spectral_dict[key]['freq_axis'][i]/1.0e6,spectral_dict[key]['spectrum'][i])
      outfile.writelines(line)
   outfile.close()


